import json
import sys
import datetime
import time

import boto3

client = boto3.client('stepfunctions')


def run(stateMachineArn, args):
    dates = args.get("dates")
    times = {}
    length = 0
    for arg in dates:
        list = []
        timeformat = arg['format']
        start = datetime.datetime.strptime(arg['start'], timeformat)
        end = datetime.datetime.strptime(arg['end'], timeformat)
        interval = int(arg["interval"])
        unit = str(arg["unit"])
        while start <= end:
            list.append(datetime.datetime.strftime(start, timeformat))
            length += 1
            if unit == 'day':
                start = start + datetime.timedelta(days=interval)
            elif unit == 'month':
                start = addmonth(start, interval)
            elif unit == 'year':
                start = addmonth(start, interval * 12)
        times[arg["name"]] = list

    length = int(length / len(dates))

    del args['dates']
    for i in range(0, length):
        for one in times:
            args[one] = times[one][i]
        resp = client.start_execution(
            stateMachineArn=stateMachineArn,
            input=json.dumps(args)
        )
        name = resp['executionArn']
        status = "RUNNING"
        while status == 'RUNNING':
            time.sleep(5)
            response = client.describe_execution(
                executionArn=name
            )
            print(response)
            status = response['status']
        print("执行完成->input:" + str(args) + ", arn:" + str(name))
        time.sleep(30)


def addmonth(date, interval):
    y = date.year
    m = date.month + interval
    d = date.day
    while m > 12 or m < 1:
        if m > 12:
            m = m - 12
            y += 1
        elif m < 1:
            m = 12 + m
            y -= 1
    return datetime.datetime(year=y, month=m, day=d)


'''
传入参数1:state-machine-arn ，如 arn:aws-cn:states:cn-northwest-1:453200521012:execution:monthly_report_business
传入参数2:参数和日期范围json字符串，如
{
	"dates": [
		{
			"name": "p_date",          //stepfunction参数名
			"start": "2022-01-01",      //开始日期
			"end": "2022-10-01",        //结束日期
			"interval": "1",            //时间间隔
			"unit": "month",            //时间间隔单位
			"format": "%Y-%m-%d"        //日期格式
		},
		{
			"name": "pDate",
			"start": "2021-12",
			"end": "2022-09",
			"interval": "1",
			"unit": "month",
			"format": "%Y-%m"
		}
	],
	"tempreture":12
}
代表，p_date从2022-01-01到2022-10-01，间隔1个月跑一次，时间格式%Y-%m-%d
     pDate从2021-12到2022-09每个一个月跑一次，时间格式%Y-%m
dates意外的参数，如tempreture:12代表每一次运行都有tempreture:12参数。
执行stepinvoke.sh
如:sh stepinvoke.sh arn:aws-cn:states:cn-northwest-1:453200521012:stateMachine:monthly_report_business "{\"dates\": [{\"name\": \"p_date\",\"start\": \"2022-02-01\",\"end\": \"2022-10-01\",\"interval\": \"1\",\"unit\": \"month\",\"format\": \"%Y-%m-%d\"},{\"name\": \"pDate\",\"start\": \"2022-01\",\"end\": \"2022-09\",\"interval\": \"1\",\"unit\": \"month\",\"format\": \"%Y-%m\"}]}"
'''
if __name__ == '__main__':
    event = sys.argv
    stateMachineArn = sys.argv[1]
    args = json.loads(sys.argv[2])
    run(stateMachineArn, args)
    print("任务完成")
